import pandas as pd
import numpy as np
import logging
import torch
from datetime import datetime
from typing import Dict, List, Tuple, Optional, Any, Union
from pathlib import Path

# 导入ART相关模块
from ..unified_representation.representaion import extract_sld_ild_tensors
from ..system_feature import SystemFeatureVector

logger = logging.getLogger(__name__)

class FeatureExtractor:
    """
    特征提取类，负责提取SLD和ILD特征并计算相似度
    """
    
    def __init__(self, output_dir: Optional[Path] = None):
        """
        初始化特征提取器
        
        参数:
            output_dir: 特征输出路径
        """
        self.output_dir = output_dir
        self.feature_dir = self.output_dir / "features" if self.output_dir else None
        if self.feature_dir:
            self.feature_dir.mkdir(exist_ok=True, parents=True)
    
    def extract_sld_ild_features(
        self,
        model,
        test_samples: List[Tuple],
        idx_to_cmdb_id: Dict[int, str],
        top_k_value: int = 3,
        anomaly_id: str = None,
        output_dir: Optional[str] = None,
        auto_save: bool = False
    ) -> SystemFeatureVector:
        """
        同时提取SLD和ILD特征，使用新的基于L1范数加权的方法
        
        参数:
            model: 训练好的模型
            test_samples: 测试样本
            idx_to_cmdb_id: 索引到CMDB ID的映射
            top_k_value: 选择的top-k实例数量
            anomaly_id: 异常ID
            output_dir: 输出目录路径
            auto_save: 是否自动保存特征和结果
            
        返回:
            SystemFeatureVector: 包含特征和相似度结果的对象
        """
        logger.info(f"开始提取SLD和ILD特征 (top_k={top_k_value})")
        
        # 获取时间戳
        timestamps = []
        for batch_ts, _, _, _ in test_samples:
            if isinstance(batch_ts, (int, float)):
                timestamps.append(batch_ts)
            else:
                timestamps.extend(batch_ts)
        
        # 提取特征张量
        sld_tensor, ild_tensor = extract_sld_ild_tensors(
            model, test_samples, top_k_value
        )
        
        # 确保时间戳长度匹配
        timestamps = timestamps[:len(sld_tensor)]
        
        # 创建并返回SystemFeatureVector对象
        feature_vector = SystemFeatureVector(
            sld_tensor=sld_tensor,
            ild_tensor=ild_tensor,
            timestamps=timestamps,
            idx_to_cmdb_id=idx_to_cmdb_id,
            anomaly_id=anomaly_id,
            output_dir=output_dir,
            auto_save=auto_save
        )
        
        return feature_vector
